package org.jflame.uid.worker.dao.impl;

import org.jflame.uid.worker.dao.IWorkerNodeDao;
import org.jflame.uid.worker.entity.WorkerNodeEntity;

import java.util.Arrays;

import org.springframework.data.redis.connection.RedisConnectionFactory;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;

public class WorkerNodeRedisImpl implements IWorkerNodeDao {

    private final String root_key = "jflame:workeridroot";
    private StringRedisTemplate redisTemplate;

    public WorkerNodeRedisImpl(RedisConnectionFactory connectionFactory) {
        redisTemplate = new StringRedisTemplate(connectionFactory);
    }

    @Override
    public WorkerNodeEntity getWorkerNodeByHostPort(String host, String port) {
        String value = buildNodeKey(host, port);
        Double score = redisTemplate.opsForZSet()
                .score(root_key, value);
        if (score != null) {
            WorkerNodeEntity node = new WorkerNodeEntity();
            node.setHostName(host);
            node.setPort(port);
            node.setId(score.intValue());
            return node;
        }
        return null;
    }

    @Override
    public boolean registerWorkerNode(WorkerNodeEntity node) {
        RedisScript<Long> luascript = new DefaultRedisScript<>(registerScript(), Long.class);
        Long scoreDouble = redisTemplate.execute(luascript, Arrays.asList(root_key),
                buildNodeKey(node.getHostName(), node.getPort()));
        node.setId(scoreDouble.intValue());
        return true;
    }

    private String buildNodeKey(String host, String port) {
        return host + "_" + port;
    }

    private String registerScript() {
        StringBuilder sb = new StringBuilder();
        sb.append("if redis.call('EXISTS',KEYS[1])==0 then ");
        sb.append("   redis.call('ZADD',KEYS[1],1,ARGV[1]) ");
        sb.append("   return 1 ");
        sb.append("end ");
        sb.append("local wokerscore=redis.call('ZSCORE',KEYS[1],ARGV[1]) ");
        sb.append("if wokerscore ~=false then ");
        sb.append("   return wokerscore ");
        sb.append("end ");
        sb.append("local maxscores=redis.call('ZRANGE',KEYS[1],-1,-1,'WITHSCORES') ");
        sb.append("if maxscores[1]~=nil then ");
        sb.append("   local wokerscore=maxscores[2]+1 ");
        sb.append("   redis.call('ZADD',KEYS[1],wokerscore,ARGV[1]) ");
        sb.append("   return wokerscore ");
        sb.append("else ");
        sb.append("   redis.call('ZADD',KEYS[1],1,ARGV[1]) ");
        sb.append("   return 1 ");
        sb.append("end ");
        return sb.toString();
    }

}
